package com.hanlin.tree;

/**
 * 描述：红黑树
 *
 * @author:hl.yuan
 * @date:2023-04-22
 */
public class RBTree<K extends Comparable<K>,V> {

    /**
     * 红黑树 结构
     *
     * @param <K>
     * @param <V>
     */
    static class RBNode<K extends Comparable<K>, V> {
        /**
         * 父节点
         */
        private RBNode parent;
        /**
         * 左节点
         */
        private RBNode left;
        /**
         * 右节点
         */
        private RBNode right;
        /**
         * 颜色
         */
        private boolean color;
        /**
         * key
         */
        private K key;
        /**
         * value值
         */
        private V value;


        public RBNode() {
        }

        public RBNode(K key, V value) {
            this.key = key;
            this.value = value;
        }

        public RBNode(K key, V value,RBNode parent) {
            this.parent = parent;
            this.key = key;
            this.value = value;
        }

        public RBNode getParent() {
            return parent;
        }

        public void setParent(RBNode parent) {
            this.parent = parent;
        }

        public RBNode getLeft() {
            return left;
        }

        public void setLeft(RBNode left) {
            this.left = left;
        }

        public RBNode getRight() {
            return right;
        }

        public void setRight(RBNode right) {
            this.right = right;
        }

        public boolean isColor() {
            return color;
        }

        public void setColor(boolean color) {
            this.color = color;
        }

        public K getKey() {
            return key;
        }

        public void setKey(K key) {
            this.key = key;
        }

        public V getValue() {
            return value;
        }

        public void setValue(V value) {
            this.value = value;
        }
    }

    private static final boolean RED = false;
    private static final boolean BLACK = true;

    // 红黑树的 root 节点
    private RBNode root;

    public RBNode getRoot() {
        return root;
    }

    public void setRoot(RBNode root) {
        this.root = root;
    }

    /**
     *          4                 6
     *         /\                /\
     *        3  6     =>       4 7
     *          /\             /\
     *          5 7           3  5
     * 左旋
     * 1. 5的父节点指向4
     * 2. 6变成root节点，
     * 3. 4的父节点指向6，
     * 4. 6的左节点指向4
     *
     * @param p
     */
    private void leftRotate(RBNode p) {
        // 边界
        if (p == null) {
            return;
        }
        // 获取右侧节点
        RBNode r = p.right;
        if (r.left != null) {
            r.left.parent = p; // 5的父节点指向4
        }
        r.parent = p.parent;// 不管p是否存在父节点，我们都设置p的父节点也为 pr的父节点
        // 如果P的父节点为空，说明P就是ROOT节点，这里左旋，就需要把root设置成功P.right
        if (p.parent == null) {
            this.root = r;
        } else if (p.parent.left == p) {
            // 如果P的父节点的左侧是自己，那么说明P是左侧节点，就需要把R换到左侧
            p.parent.left = r;
        } else {
            p.parent.right = r;
        }
        // 6的左节点指向4
        r.left = p;
        // 4的父节点指向6
        p.parent = r;
    }

    /**
     *       6                   4
     *      /\                  /\
     *     4  7        =>      3  6
     *    /\                     /\
     *   3  5                   5  7
     * 右旋
     * 1. 5的父节点指向6
     * 2. 4变成root节点
     * 3. 6的父节点指向4
     * 4. 4的左节点指向6
     * @param p
     */
    private void rightRotate(RBNode p) {
        // 边界
        if (p == null) {
            return;
        }
        // 获取左侧节点
        RBNode r = p.left;
        if (r.right != null) {
            r.right.parent = p;  // 相当于5的父节点指向6
        }
        r.parent = p.parent;
        // 如果P的父节点为空，说明P就是ROOT节点，这里左旋，就需要把root设置成功P.right
        if (p.parent == null) {
            this.root = r;
        } else if (p.parent.right == p) {
            // 如果P的父节点的左侧是自己，那么说明P是左侧节点，就需要把R换到左侧
            p.parent.right = r;
        } else {
            p.parent.left = r;
        }
        // 4的右节点指向6
        r.right = p;
        // 6的父节点指向4
        p.parent = r;
    }

    /**
     * 红黑树 新增节点的操作
     *
     * @param key
     * @param value
     */
    public void put(K key, V value) {
        RBNode t = root;
        // 如果root节点为空，就可以直接插入
        if (t == null) {
            // 这是插入的第一个节点
            root = new RBNode(key, value == null ? key : value, null);
            return;
        }
        int cmp;
        // 1. 找到插入的位置(找到新增节点的父节点)
        RBNode parent;
        do {
            parent = t;
            cmp = key.compareTo((K) t.key); // 主要是判断当前节点插入的位置，看是插入在左侧还是在右侧
            if (cmp < 0) { // 小于放在左侧
                // 从左侧查找
                t = t.left;
            } else if (cmp > 0) { // 大于放在右侧
                // 从右侧查找
                t = t.right;
            } else { // 等于，直接覆盖掉
                // 说明节点存在，用插入节点的值覆盖掉 相同的节点
                t.setValue(value == null ? key : value);
                return;
            }
        } while (t != null); // 这里必选要循环去比较，因为不知道要比较多少次
        // 2.将新节点添加到父节点的子节点中(左右子节点)
        // 创建要插入的节点
        RBNode node = new RBNode(key, value == null ? key : value, parent);
        if (cmp < 0) {
            // 将新节点添加到父节点的左侧
            parent.left = node;
        } else {
            // 添加到右侧
            parent.right = node;
        }
        // 旋转和变色 调整红黑树的平衡
        fixAfterPut(node);
    }

    private RBNode parentOf(RBNode node){
        return node != null ? node.parent :null;
    }

    private  RBNode leftOf(RBNode node){
        return node != null ? node.left:null;
    }

    private  RBNode rightOf(RBNode node){
        return node != null ? node.right :null;
    }

    private boolean colorOf(RBNode node){
        return node == null ? BLACK:node.color;
    }

    private void setColor(RBNode node,boolean color){
        if(node != null){
            node.setColor(color);
        }
    }

    /**
     * 插入节点后的调整操作
     * 2-3-4对应的操作
     *   2节点： 新插入一个元素 直接和2节点合并 不用调整
     *       红黑树：新增一个红色节点在黑色节点下 不需要调整
     *   3节点： 新插入一个元素在3节点下 那么会出现6中情况(2两种不需要调整，4中需要调整)
     *       红黑树： 插入的节点是 插入在 上黑下红的结构中，插入在红色节点
     *   4节点： 新插入一个元素在4节点下，那么会出现4中情况 都需要调整
     *       红黑树：新增的节点是红色 爷爷节点是黑色，父亲节点和叔叔节点是红色
     * @param x
     */
    private void fixAfterPut(RBNode<K,Object> x){
        // 插入的节点 肯定红色
        x.color = RED;
        // 2节点不用调整，3,4节点才需要调整
        while( x != null && x != root && x.parent.color == RED){
            if(parentOf(x) == leftOf(parentOf(parentOf(x)))){ // 当前节点的父节点  = 当前的节点的爷爷的左侧节点，具体看【红黑树_变色穷举.png】 左侧
                // 需要调整的只剩下4种情况，有叔叔节点和没有叔叔节点
                // 获取插入节点的叔叔节点
                RBNode y = rightOf(parentOf(parentOf(x)));
                if(colorOf(y) == RED){
                    // 说明是有叔叔节点的
                    // 变色+递归
                    // 父亲和叔叔变为黑色  爷爷变为红色
                    setColor(parentOf(x),BLACK);
                    setColor(y,BLACK);
                    setColor(parentOf(parentOf(x)),RED);
                    // 递归处理
                    x = parentOf(parentOf(x));
                }else{
                    // 说明是没有叔叔节点的
                    // 如果当前节点是父节点的右侧：这里为什么要判断，是因为这里需要再进行一次左旋，具体看 【红黑树_变色穷举.png】 左侧第二图
                    if(x == parentOf(x).right){
                        x = parentOf(x);
                        leftRotate(x); // 左旋
                    }
                    // 父亲节点变为黑色 爷爷节点变为红色
                    setColor(parentOf(x),BLACK);
                    setColor(parentOf(parentOf(x)),RED);
                    // 根据爷爷节点做右旋操作
                    rightRotate(parentOf(parentOf(x)));
                }
            }else{
                // 需要调整的也是4中情况，刚好和上面的4中情况左右相反
                // 需要调整的只剩下4种情况，有叔叔节点和没有叔叔节点
                // 获取插入节点的叔叔节点
                RBNode y = leftOf(parentOf(parentOf(x)));
                if(colorOf(y) == RED){
                    // 说明是有叔叔节点的
                    // 变色+递归
                    // 父亲和叔叔变为黑色  爷爷变为红色
                    setColor(parentOf(x),BLACK);
                    setColor(y,BLACK);
                    setColor(parentOf(parentOf(x)),RED);
                    // 递归处理
                    x = parentOf(parentOf(x));
                }else{
                    // 说明是没有叔叔节点的
                    if(x == parentOf(x).left){
                        x = parentOf(x);
                        rightOf(x);
                    }
                    // 父亲节点变为黑色 爷爷节点变为红色
                    setColor(parentOf(x),BLACK);
                    setColor(parentOf(parentOf(x)),RED);
                    // 根据爷爷节点做右旋操作
                    leftRotate(parentOf(parentOf(x)));
                }
            }
        }
        // root 节点肯定为黑色
        root.color = BLACK;
    }


}
